module TVD_mod

  use constants_mod
  use namelist_mod
  use boundary_mod, only : set_periodic_boundary, pbc_idx
  
  implicit none

  private
  public compute_interface_flux_TVD
  public compute_flux_jacobian_TVD

  contains

  subroutine compute_interface_flux_TVD(q, dt, flux)

    real(r8), intent(in) :: q(1-nhalo:nx+nhalo)
    real(r8), intent(in) :: dt
    real(r8), intent(out) :: flux(0:nx)
    real(r8) alpha
    integer i
    real(r8) gradq, phi, kesi, slope(1-nhalo:nx+nhalo), slope1, slope2
    logical method1, method2

    method1 = .true.
    method2 = .false.

    kesi = sign(1.0_r8, u)
    !===============Method 1
    if (method1) then
    do i = 0, nx
      ! the  ratio of upwind-side gradient to downwind-side gradient (r) determines the value of function Phi
      gradq = merge((q(i  ) - q(i-1)) / (q(i+1) - q(i) + epsilon), &
                    (q(i+2) - q(i+1)) / (q(i+1) - q(i) + epsilon), u >= 0.0_r8) ! eq. (3.47) in Chapter 4
      if (flux_limiter == 'minmod') then
        phi = minmod(1.0_r8, gradq) ! or max(0.0_r8, min(1.0_r8, gradq))
      else if (flux_limiter == 'superbee') then
        phi = max(0.0_r8, min(1.0_r8, 2.0_r8 * gradq), min(2.0_r8, gradq))
      else if (flux_limiter == 'vanleer') then
        phi = (gradq + abs(gradq)) / (1.0_r8 + abs(gradq))
      else if (flux_limiter == 'MC') then
        phi = max(0.0_r8, min((1.0_r8 + gradq) / 2.0_r8, 2.0_r8, 2.0_r8 * gradq))
      end if
      ! based on upwind and Lax-Wendroff scheme
      flux(i) = u * (0.5 * ((1.0_r8 + kesi) * q(i) + (1.0_r8 - kesi) * q(i+1)) + &
               phi * (0.5_r8 * (kesi - u * dt / dx) * (q(i+1) - q(i)))) ! 等价于 eq. (4.38) in Chapter 4
    end do
    end if

    !=======================Method 2 from Chapter 4,Advection algorithms II. Flux conservation, subgrid models and flux limiters
    if (method2) then
    do i = 1, nx
      if (slope_limiter == 'minmod') then ! eq. (4.28) in Chapter 4
        slope(i) = minmod((q(i) - q(i-1)) / dx, (q(i+1) - q(i)) / dx)
      else if (slope_limiter == 'superbee') then ! eq. (4.30) in Chapter 4
        slope1 = minmod((q(i+1) - q(i)) / dx, 2.0_r8 * (q(i) - q(i-1)) / dx)
        slope2 = minmod(2.0_r8 * (q(i+1) - q(i)) / dx, (q(i) - q(i-1)) / dx)
        slope(i) = maxmod(slope1, slope2)
      else if (slope_limiter == 'downwind') then ! Lax-Wendroff eq. (4.24) in Chapter 4
        slope(i) = (q(i+1) - q(i)) / dx
      else if (slope_limiter == 'upwind') then ! Beam-Warming eq. (4.23) in Chapter 4
        slope(i) = (q(i) - q(i-1)) / dx
      else if (slope_limiter == 'centered') then ! Fromm' method eq. (4.22) in Chapter 4
        slope(i) = (q(i+1) - q(i-1)) / (2.0_r8 * dx)
      else if (slope_limiter == 'zero') then ! doner-cell 
        slope(i) = 0.0_r8
      end if
    end do
    call set_periodic_boundary(slope)
    do i = 0, nx
      ! eq. (4.35) in Chapter 4 ! piecewise linear scheme assumption
      flux(i) = 0.5 * u * ((1.0_r8 + kesi) * q(i) + (1.0_r8 - kesi) * q(i+1)) + &
              0.25_r8 * abs(u) * (1.0_r8 - abs(u * dt / dx)) * dx * &
            ((1.0_r8 + kesi) * slope(i) + (1.0_r8 - kesi) * slope(i+1))
    end do
    end if

  end subroutine compute_interface_flux_TVD

  subroutine compute_flux_jacobian_TVD(q, dt, jac)

    real(r8), intent(in) :: q(1-nhalo:nx+nhalo)
    real(r8), intent(in) :: dt
    real(r8), intent(out) :: jac(1:nx, 1:nx)
    integer i, j, idx_m1, idx_m2, idx_p1, idx_p2
    real(r8) gradq, phi, dphi_dr, dphi_dq

    jac = 0.0_r8
    
    if (u > 0.0_r8) then
      do i = 1, nx
        ! dF_i+1/2/dq
        gradq = (q(i) - q(i-1)) / (q(i+1) - q(i) + 2 * epsilon)
        phi = (gradq + abs(gradq)) / (1.0_r8 + abs(gradq)) ! vanleer
        idx_m1 = pbc_idx(i-1, nx)
        idx_p1 = pbc_idx(i+1, nx)
        
        if (gradq > 0) then
          dphi_dr = 2.0_r8 / (1.0_r8 + gradq)**2
          dphi_dq = dphi_dr * (-1.0_r8 / (q(i+1) - q(i) + epsilon))
          jac(i, idx_m1) = 0.5_r8 * u * dphi_dq * (1.0_r8 - u * dt / dx) * (q(i+1) - q(i))

          dphi_dq = dphi_dr * (q(i+1) - q(i-1)) / (q(i+1) - q(i) + epsilon)**2
          jac(i, i) = u + 0.5_r8 * u * dphi_dq * (1.0_r8 - u * dt / dx) * (q(i+1) - q(i)) - &
                          0.5_r8 * u * phi * (1.0_r8 - u * dt / dx)

          dphi_dq = dphi_dr * (-(q(i) - q(i-1)) / (q(i+1) - q(i) + epsilon)**2)
          jac(i, idx_p1) = 0.5_r8 * u * dphi_dq * (1.0_r8 - u * dt / dx) * (q(i+1) - q(i)) + &
                           0.5_r8 * u * phi * (1.0_r8 - u * dt / dx)           
        end if
      end do

      do i = 1, nx
        !  DF_i-1/2/ dq 
        gradq = (q(i-1) - q(i-2)) / (q(i) - q(i-1) + 2 * epsilon) 
        phi = (gradq + abs(gradq)) / (1.0_r8 + abs(gradq)) ! vanleer
        idx_m1 = pbc_idx(i-1, nx)
        idx_m2 = pbc_idx(i-2, nx)
        if (gradq > 0) then
          dphi_dr = 2.0_r8 / (1.0_r8 + gradq)**2
          dphi_dq = dphi_dr * (-1.0_r8 / (q(i) - q(i-1) + epsilon))
          jac(i, idx_m2) = jac(i, idx_m2) - (0.5_r8 * u * dphi_dq * (1.0_r8 - u * dt / dx) * (q(i) - q(i-1)))

          dphi_dq = dphi_dr * ((q(i) - q(i-2)) / (q(i) - q(i-1) + epsilon)**2)
          jac(i, idx_m1) = jac(i, idx_m1) - (u + 0.5_r8 * u * dphi_dq * (1.0_r8 - u * dt / dx) * (q(i) - q(i-1)) - &
                                                 0.5_r8 * u * phi * (1.0_r8 - u * dt / dx))

          dphi_dq = dphi_dr * (-(q(i-1) - q(i-2)) / (q(i) - q(i-1) + epsilon)**2)
          jac(i, i) = jac(i, i) - (0.5_r8 * u * dphi_dq * (1.0_r8 - u * dt / dx) * (q(i) - q(i-1)) + &
                                   0.5_r8 * u * phi * (1.0_r8 - u * dt / dx))
        end if
      end do
    else
      do i = 1, nx
        ! dF_i+1/2/ dq 
        gradq = (q(i+2) - q(i+1)) / (q(i+1) - q(i) + epsilon)
        phi = (gradq + abs(gradq)) / (1.0_r8 + abs(gradq))
        idx_p1 = pbc_idx(i+1, nx)
        idx_p2 = pbc_idx(i+2, nx)
        if (gradq > 0.0_r8) then
          dphi_dr = 2.0_r8 / (1.0_r8 + gradq)**2
          dphi_dq = dphi_dr * ((q(i+2) - q(i+1)) / (q(i+1) - q(i) + epsilon)**2)
          jac(i,i) = 0.5_r8 * u * dphi_dq * (-1.0_r8 - u * dt / dx) * (q(i+1) - q(i)) - &
                     0.5_r8 * u * phi * (-1.0_r8 - u * dt / dx)
          
          dphi_dq = dphi_dr * ((q(i) - q(i+2)) / (q(i+1) - q(i) + epsilon)**2)
          jac(i, idx_p1) = u + 0.5_r8 * u * dphi_dq * (-1.0_r8 - u * dt / dx) * (q(i+1) - q(i)) + &
                               0.5_r8 * u * phi * (-1.0_r8 - u * dt / dx)

          dphi_dq = dphi_dr / (q(i+1) - q(i) + epsilon)
          jac(i, idx_p2) = 0.5_r8 * u * dphi_dq * (-1.0_r8 - u * dt / dx) * (q(i+1) - q(i))
        end if

        !  DF_i-1/2/ dq 
        gradq = (q(i+1) - q(i)) / (q(i) - q(i-1) + epsilon)
        phi = (gradq + abs(gradq)) / (1.0_r8 + abs(gradq))
        idx_m1 = pbc_idx(i-1, nx)
        idx_p1 = pbc_idx(i+1, nx)
        if (gradq > 0.0_r8) then
          dphi_dr = 2.0_r8 / (1.0_r8 + gradq)**2
          dphi_dq = dphi_dr * ((q(i+1) - q(i)) / (q(i) - q(i-1) + epsilon)**2)
          jac(i, idx_m1) = jac(i, idx_m1) - (0.5_r8 * u * dphi_dq * (-1.0_r8 - u * dt / dx) * (q(i) - q(i-1)) - &
                                             0.5_r8 * u * phi * (-1.0_r8 - u * dt / dx))
          
          dphi_dq = dphi_dr * ((q(i-1) - q(i+1)) / (q(i) - q(i-1) + epsilon)**2)
          jac(i, i) = jac(i,i) - (u + 0.5_r8 * u * dphi_dq * (-1.0_r8 - u * dt / dx) * (q(i) - q(i-1)) + &
                                      0.5_r8 * u * phi * (-1.0_r8 - u * dt / dx))
          
          dphi_dq = dphi_dr / (q(i) - q(i-1) + epsilon)
          jac(i, idx_p1) = jac(i, idx_p1) - 0.5_r8 * u * dphi_dq * (-1.0_r8 - u * dt / dx) * (q(i) - q(i-1))
        end if
      end do
    end if
    jac = -jac / dx

  end subroutine compute_flux_jacobian_TVD

  function minmod(a, b)

    real(r8), intent(in) :: a, b
    real(r8) minmod
    
    if (abs(a) < abs(b) .and. a * b > 0.0_r8) then
      minmod = a
    elseif (abs(a) > abs(b) .and. a * b > 0.0_r8) then
      minmod = b
    else
      minmod = 0.0_r8
    end if

  end function minmod
  
  function maxmod(a, b)

    real(r8), intent(in) :: a, b
    real(r8) maxmod

    if (abs(a) > abs(b) .and. a * b > 0.0_r8) then
      maxmod = a
    elseif (abs(a) < abs(b) .and. a * b > 0.0_r8) then
      maxmod = b
    else
      maxmod = 0.0_r8
    end if

  end function maxmod

end module TVD_mod